#!/usr/bin/env python
# -*- coding: UTF-8 -*-
'''
@Project ：V2 
@File    ：gas_pipeline_qc_partitio1.py
@IDE     ：PyCharm 
@Author  ：郭星
@Date    ：2025/9/23 14:52 
'''
#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Gas Mixing -> QC Scheduling (component-wise MIX; device can test up to 2 components)
------------------------------------------------------------------------------------
- MIX: same-recipe bottles can be batched; charging is done component-by-component
       with sequence-dependent setups (by last component), even for a batch.
- QC : each device can cover up to K components (here K<=2) as its capability set.
       For each bottle we decide a grouping of its required components:
         * size==1  -> serial on one device (single component)
         * size>=2  -> 1-pass: occupy several devices simultaneously, same [st,en];
                       each device can test up to 'QC_DEVICE_CAPACITY_PER_PASS'
                       components simultaneously; its per-device time aggregates
                       multiple-component times by 'QC_DEVICE_AGG_TIME' = max|sum.
       No mid-detach during 1-pass. Start time = max(mix_end, all devices ready+setup).
- GA  : searches per-bottle grouping policy (ALL, SPLIT_LONGEST, EACH), minimizing Cmax.

Author: you
"""

import random
import numpy as np
from dataclasses import dataclass
from typing import List, Dict, Tuple, Set, Optional

try:
    import matplotlib.pyplot as plt
    HAS_MPL = True
except Exception:
    HAS_MPL = False


# ============================ CONFIG ============================ #

CONFIG = {
    # ---- scale ----
    "NUM_BOTTLES": 20,
    "NUM_MIX_MACHINES": 2,
    "NUM_QC_DEVICES": 6,
    "MAX_COMPONENTS_PER_BOTTLE": 4,

    # ---- MIX (charging) ----
    "MIX_BATCH_CAP": 300,                 # max bottles per batch (same recipe)
    "MIX_RATE_PER_H": [180.0, 150.0],     # bottles/hour per mixer (not used if using per-bottle times)
    "MIX_SETUP_SAME": 0.05,               # setup if last_component == current_component
    "MIX_SETUP_DIFF": 0.35,               # setup otherwise
    # per-bottle charging time (hours) for each component
    "COMPONENT_FILL_TIME": {
        "O2": 0.015, "CO2": 0.020, "CH4": 0.022, "CO": 0.018,
        "H2": 0.012, "N2": 0.010, "C2H6": 0.025
    },
    # charging component order priority (left earlier); others by alpha
    "CHARGE_PRIORITY": ["CO2", "O2", "CH4", "H2", "CO", "N2", "C2H6"],

    # ---- QC (testing) ----
    "COMPONENT_TYPES": ["O2", "CO2", "CH4", "CO", "H2", "N2", "C2H6"],
    "COMPONENT_TEST_TIME": {              # single-device time per bottle
        "O2": 0.18, "CO2": 0.22, "CH4": 0.25, "CO": 0.20,
        "H2": 0.15, "N2": 0.12, "C2H6": 0.28
    },
    "QC_SETUP_TIME": 0.12,                # change of method signature (set of components)

    # ---- Device capability & 1-pass capacity ----
    "QC_DEVICE_CAPACITY_PER_PASS": 2,     # a device can test up to K components simultaneously in a 1-pass
    "QC_DEVICE_AGG_TIME": "max",          # per-device time when testing multiple comps at once: "max" or "sum"

    # ---- GA over QC grouping policy ----
    "GA_POP": 80,
    "GA_GENS": 200,
    "GA_PC": 0.9,
    "GA_PM": 0.35,
    "GA_ELITE": 10,

    # ---- plotting ----
    "LABEL_WRAP_CHARS": 22,               # wrap text width; all texts in black
    "RANDOM_SEED": 7
}


# ============================ Data Models ============================ #

@dataclass(frozen=True)
class Recipe:
    components: Tuple[Tuple[str, float], ...]  # (component, concentration%)

    def key(self) -> Tuple[Tuple[str,float], ...]:
        # group by identical components+concentrations for batching
        return tuple(sorted(self.components, key=lambda x: x[0]))


@dataclass
class Bottle:
    id: int
    recipe: Recipe


@dataclass
class MixTask:
    start: float
    end: float
    machine: int
    label: str


@dataclass
class QCTask:
    start: float
    end: float
    device: int
    label: str
    group_key: Optional[Tuple[int,float,float,int]] = None  # (bid, st, en, group_index) for 1-pass coloring


# ============================ Random Instance ============================ #

class RandomData:
    """
    - Each QC device has capability set of 1~2 components (drawn from COMPONENT_TYPES).
    - Each bottle has 1..K components (drawn from union of device capabilities).
    """
    def __init__(self, cfg):
        self.cfg = cfg
        random.seed(cfg["RANDOM_SEED"]); np.random.seed(cfg["RANDOM_SEED"])

    def gen(self):
        N = self.cfg["NUM_BOTTLES"]
        D = self.cfg["NUM_QC_DEVICES"]
        types = self.cfg["COMPONENT_TYPES"]
        Kmax = self.cfg["MAX_COMPONENTS_PER_BOTTLE"]

        # device capability sets (size 1 or 2)
        dev_caps: List[Set[str]] = []
        for i in range(D):
            if len(types) >= 2:
                k = random.choice([1, 2])
            else:
                k = 1
            dev_caps.append(set(random.sample(types, k)))

        union_caps = set().union(*dev_caps) if dev_caps else set(types)
        if not union_caps:
            union_caps = set(types)

        bottles: List[Bottle] = []
        for i in range(N):
            k = np.random.randint(1, Kmax+1)
            k = min(k, len(union_caps))
            comps = random.sample(list(union_caps), k)
            comps_t = tuple((c, round(np.random.uniform(0.1, 5.0), 3)) for c in comps)
            bottles.append(Bottle(i, Recipe(comps_t)))
        return bottles, dev_caps


# ============================ Scheduler ============================ #

class Scheduler:
    """
    MIX: component-wise charging with batching and component-based setups.
    QC : per-bottle grouping; groups of size>=2 are 1-pass with simultaneous devices.
         One device can test up to K comps simultaneously in a 1-pass, per-device time
         aggregated by "max" or "sum". Signature = frozenset(components tested on device).
    """
    def __init__(self, cfg, bottles: List[Bottle], dev_caps: List[Set[str]]):
        self.cfg = cfg
        self.bottles = bottles
        self.N = len(bottles)
        self.dev_caps = dev_caps
        self.D_mix = cfg["NUM_MIX_MACHINES"]
        self.D_qc  = cfg["NUM_QC_DEVICES"]

        rates = cfg["MIX_RATE_PER_H"]
        if len(rates) < self.D_mix:
            mul = (self.D_mix + len(rates) - 1) // len(rates)
            rates = (rates * mul)[:self.D_mix]
        self.mix_rates = rates

    # ---------- charging component order ----------
    def charging_sequence(self, comps: List[str]) -> List[str]:
        pri = {c: idx for idx, c in enumerate(self.cfg["CHARGE_PRIORITY"])}
        return sorted(comps, key=lambda c: (pri.get(c, 10_000), c))

    # ---------- MIX: component-wise charging with batching ----------
    def schedule_mixing(self) -> Tuple[Dict[int,float], List[MixTask]]:
        cap = self.cfg["MIX_BATCH_CAP"]
        setup_same = self.cfg["MIX_SETUP_SAME"]
        setup_diff = self.cfg["MIX_SETUP_DIFF"]
        fill_t = self.cfg["COMPONENT_FILL_TIME"]

        # group bottles by recipe key
        key2ids: Dict[Tuple, List[int]] = {}
        for b in self.bottles:
            key2ids.setdefault(b.recipe.key(), []).append(b.id)

        # split each group into batches (<=cap)
        batches: List[List[int]] = []
        for _, ids in key2ids.items():
            cur=[]
            for bid in ids:
                cur.append(bid)
                if len(cur) >= cap:
                    batches.append(cur); cur=[]
            if cur: batches.append(cur)

        # stable scheduling priority: larger batches first
        batches.sort(key=lambda bl: -len(bl))

        ready = [0.0]*self.D_mix
        last_comp: List[Optional[str]] = [None]*self.D_mix
        mix_tasks: List[MixTask] = []
        mix_end: Dict[int,float] = {}

        for bl in batches:
            comps = [c for c,_ in self.bottles[bl[0]].recipe.components]
            seq = self.charging_sequence(comps)
            batch_prev_end = 0.0
            for comp in seq:
                t_per_bottle = float(fill_t.get(comp, 0.02))
                pt = t_per_bottle * len(bl)  # per-bottle charging time * batch size

                best=None
                for m in range(self.D_mix):
                    stp = (
                        (setup_same if (last_comp[m] == comp and last_comp[m] is not None) else
                         (0.0 if last_comp[m] is None else setup_diff))
                    )
                    st = max(ready[m] + stp, batch_prev_end)
                    en = st + pt
                    if (best is None) or (en < best[0]):
                        best=(en,m,st,pt,stp)
                en,m,st,pt_use,stp = best
                mix_tasks.append(MixTask(st, en, m, f"Batch({','.join('S'+str(i) for i in bl)})-{comp}"))
                ready[m] = en; last_comp[m] = comp
                batch_prev_end = en

            for bid in bl:
                mix_end[bid] = batch_prev_end

        return mix_end, mix_tasks

    # ---------- QC helpers ----------
    def single_test_time(self, comp: str) -> float:
        return float(self.cfg["COMPONENT_TEST_TIME"].get(comp, 0.2))

    def comp_to_devs(self) -> Dict[str, List[int]]:
        mapping: Dict[str, List[int]] = {}
        for d, caps in enumerate(self.dev_caps):
            for c in caps:
                mapping.setdefault(c, []).append(d)
        return mapping

    # policy codes: 0=ALL_TOGETHER, 1=SPLIT_LONGEST, 2=EACH_SEPARATE
    def make_groups(self, comps: List[str], policy_code: int) -> List[List[str]]:
        if len(comps) <= 1:
            return [comps[:]]
        times = {c: self.single_test_time(c) for c in comps}
        if policy_code == 0:
            return [comps[:]]
        if policy_code == 1:
            longest = max(comps, key=lambda c: times[c])
            rest = [c for c in comps if c != longest]
            return [[longest], rest] if rest else [[longest]]
        # EACH_SEPARATE; order by preferred device id then by name
        c2d = self.comp_to_devs()
        comps_sorted = sorted(comps, key=lambda c: (min(c2d.get(c,[999])) if c2d.get(c) else 999, c))
        return [[c] for c in comps_sorted]

    # ---------- QC scheduling with groups ----------
    def schedule_qc_with_groups(self,
                                mix_end: Dict[int,float],
                                bottle_groups: Dict[int, List[List[str]]]
                               ) -> Tuple[float, List[QCTask]]:
        D = self.D_qc
        ready = [0.0]*D
        last_sig: List[Optional[frozenset]] = [None]*D
        setup_t = self.cfg["QC_SETUP_TIME"]
        cap_per_pass = int(self.cfg.get("QC_DEVICE_CAPACITY_PER_PASS", 1))
        agg_mode = str(self.cfg.get("QC_DEVICE_AGG_TIME", "max")).lower()

        def agg_time(comps: List[str]) -> float:
            ts = [self.single_test_time(c) for c in comps]
            return max(ts) if agg_mode == "max" else sum(ts)

        comp2devs = self.comp_to_devs()
        tasks: List[QCTask] = []
        order_seq = sorted(self.bottles, key=lambda b: mix_end[b.id])

        for b in order_seq:
            bid = b.id
            t_cur = mix_end[bid]
            groups = bottle_groups[bid]  # ordered list of component lists

            for gi, group in enumerate(groups):
                if len(group) == 1:
                    # ---- single component: serial ----
                    comp = group[0]
                    cands = comp2devs.get(comp, [])
                    if not cands:
                        raise RuntimeError(f"No QC device can test component {comp}")
                    best_d=None; best_ready=None; best_stp=0.0
                    for d in cands:
                        sig = frozenset([comp])
                        stp = 0.0 if (last_sig[d] is None or last_sig[d]==sig) else setup_t
                        rdy = ready[d] + stp
                        if (best_ready is None) or (rdy < best_ready):
                            best_d, best_ready, best_stp = d, rdy, stp
                    st = max(t_cur, best_ready)
                    pt = self.single_test_time(comp)
                    en = st + pt
                    tasks.append(QCTask(st, en, best_d, f"S{bid} serial E{best_d+1}({comp})", None))
                    ready[best_d] = en
                    last_sig[best_d] = frozenset([comp])
                    t_cur = en
                    continue

                # ---- multi-component group: 1-pass ----
                remaining = set(group)
                chosen: List[Tuple[int, List[str], float]] = []  # (device, comps_on_this_device, ready_with_setup_est)

                # 贪心装箱：每轮选“可覆盖剩余组分多且就绪早”的设备，装入其能测且时长更长的组分（≤cap）
                while remaining:
                    best=None
                    for d in range(D):
                        able = list(remaining & self.dev_caps[d])
                        if not able:
                            continue
                        # 估计就绪：若上次签名包含本组任何成分，视作无setup；否则需要setup_t
                        stp_est = 0.0 if (last_sig[d] is None or (last_sig[d] & remaining)) else setup_t
                        rdy = ready[d] + stp_est
                        gain = min(len(able), cap_per_pass)
                        score = (gain, -rdy)  # gain更大优先，其次就绪更早
                        if (best is None) or (score > best[0]):
                            best = (score, d, rdy, able)
                    if best is None:
                        raise RuntimeError(f"Group {group} cannot be fully covered by devices.")
                    _, d, rdy, able = best
                    # 把该设备能测的剩余组件按单机时长降序装入（≤cap）
                    able.sort(key=lambda c: self.single_test_time(c), reverse=True)
                    take = able[:cap_per_pass]
                    chosen.append((d, take, rdy))
                    remaining -= set(take)

                # 统一起点：所有 chosen 设备真正的 ready+setup 与 t_cur 的最大
                earliest = t_cur
                per_dev = []  # (d, signature, setup, per_device_time)
                for (d, comps_on_d, _) in chosen:
                    sig = frozenset(comps_on_d)
                    stp = 0.0 if (last_sig[d] is None or last_sig[d]==sig) else setup_t
                    earliest = max(earliest, ready[d] + stp)
                    per_t = agg_time(comps_on_d)
                    per_dev.append((d, sig, stp, per_t))

                st = earliest
                en = st + max(t for (_,_,_,t) in per_dev)  # 1-pass 段长取各设备内用时的最大值

                used_str = "+".join(f"E{d+1}" for (d,_,_,_) in per_dev)
                comps_str = ",".join(sorted(group))
                label = f"S{bid} 1-pass[{used_str}] comps={comps_str} max={en-st:.2f}h"
                group_key = (bid, round(st,6), round(en,6), gi)
                for (d, sig, stp, per_t) in per_dev:
                    tasks.append(QCTask(st, en, d, label, group_key))
                    ready[d] = en
                    last_sig[d] = sig
                t_cur = en

        Cmax = max((t.end for t in tasks), default=0.0)
        return Cmax, tasks


# ============================ GA for grouping policy ============================ #

class GA:
    """
    Chromosome: int array of length N (per bottle policy)
        0 = ALL_TOGETHER, 1 = SPLIT_LONGEST, 2 = EACH_SEPARATE
    Fitness: minimize Cmax (MIX fixed; QC groups determined by chromosome)
    """
    def __init__(self, cfg, scheduler: Scheduler, mix_end: Dict[int,float]):
        self.cfg = cfg
        self.S = scheduler
        self.mix_end = mix_end
        self.N = self.S.N
        random.seed(cfg["RANDOM_SEED"]); np.random.seed(cfg["RANDOM_SEED"])

    def rand_ind(self):
        return np.random.randint(0,3,size=self.N,dtype=int)

    def crossover(self, a, b):
        n=len(a); cut=np.random.randint(1,n)
        c1=np.concatenate([a[:cut], b[cut:]])
        c2=np.concatenate([b[:cut], a[cut:]])
        return c1,c2

    def mutate(self, ind, pm=0.2):
        if np.random.rand()<pm:
            i=np.random.randint(self.N)
            ind[i]=np.random.randint(0,3)
        return ind

    def ind_to_groups(self, ind) -> Dict[int, List[List[str]]]:
        mapping={}
        for i, b in enumerate(self.S.bottles):
            comps=[c for c,_ in b.recipe.components]
            mapping[b.id]=self.S.make_groups(comps, int(ind[i]))
        return mapping

    def fitness(self, ind):
        groups = self.ind_to_groups(ind)
        Cmax, _ = self.S.schedule_qc_with_groups(self.mix_end, groups)
        return Cmax

    def run(self):
        pop=self.cfg["GA_POP"]; gens=self.cfg["GA_GENS"]
        pc=self.cfg["GA_PC"]; pm=self.cfg["GA_PM"]; elite=self.cfg["GA_ELITE"]
        pool=[]
        for _ in range(pop):
            ind=self.rand_ind()
            f=self.fitness(ind)
            pool.append((ind,f))
        best=min(pool, key=lambda x:x[1]); hist=[best[1]]

        for g in range(1, gens+1):
            pool.sort(key=lambda x:x[1])
            elites=pool[:elite]
            new=[(e[0].copy(), e[1]) for e in elites]
            while len(new)<pop:
                p1=random.choice(pool)[0]
                p2=random.choice(pool)[0]
                c1,c2=p1.copy(),p2.copy()
                if np.random.rand()<pc:
                    c1,c2=self.crossover(p1,p2)
                if np.random.rand()<pm:
                    c1=self.mutate(c1)
                if np.random.rand()<pm:
                    c2=self.mutate(c2)
                new.append((c1, self.fitness(c1)))
                if len(new)<pop:
                    new.append((c2, self.fitness(c2)))
            pool=new
            cur=min(pool, key=lambda x:x[1])
            if cur[1] < best[1]:
                best=cur
            hist.append(best[1])
        return best[0], best[1], hist


# ============================ Plotting ============================ #

def wrap_text(s: str, width: int) -> str:
    out=[]; line=""
    for ch in s:
        line += ch
        if len(line) >= width and ch == ' ':
            out.append(line.rstrip()); line=""
    if line: out.append(line)
    return "\n".join(out)

def plot_gantt(mix_tasks: List[MixTask], qc_tasks: List[QCTask], cfg,
               dev_caps: List[Set[str]],
               title="Gantt: MIX (component-wise) + QC (grouped 1-pass/serial)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return

    # rows
    rows=[]
    for m in range(cfg["NUM_MIX_MACHINES"]):
        rows.append(("MIX", m))
    for d in range(cfg["NUM_QC_DEVICES"]):
        rows.append(("QC", d))
    row_idx={("MIX",m): m for m in range(cfg["NUM_MIX_MACHINES"])}
    off=cfg["NUM_MIX_MACHINES"]
    for d in range(cfg["NUM_QC_DEVICES"]):
        row_idx[("QC",d)] = off + d

    # colors for 1-pass groups
    from itertools import cycle
    palette = cycle(plt.rcParams['axes.prop_cycle'].by_key()['color'])
    grp_color: Dict[Tuple[int,float,float,int], str] = {}
    for t in qc_tasks:
        if t.group_key and t.group_key not in grp_color:
            grp_color[t.group_key] = next(palette)

    fig_h = max(4.5, 0.6*len(rows)+2)
    fig, ax = plt.subplots(figsize=(12, fig_h))

    # MIX
    for t in mix_tasks:
        y=row_idx[("MIX",t.machine)]
        ax.barh(y, t.end-t.start, left=t.start, height=0.8,
                edgecolor='black', linewidth=0.6)
        ax.text((t.start+t.end)/2, y, wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='black', fontsize=8)

    # QC
    for t in qc_tasks:
        y=row_idx[("QC",t.device)]
        color = grp_color.get(t.group_key, None)
        ax.barh(y, t.end-t.start, left=t.start, height=0.8,
                color=color, edgecolor='black', linewidth=0.6)
        ax.text((t.start+t.end)/2, y, wrap_text(t.label, cfg["LABEL_WRAP_CHARS"]),
                ha='center', va='center', color='black', fontsize=8)

    # y labels (show device capabilities)
    yticks=[]; ylabels=[]
    for m in range(cfg["NUM_MIX_MACHINES"]):
        yticks.append(row_idx[("MIX", m)])
        ylabels.append(f"MIX-M{m+1}")
    for d in range(cfg["NUM_QC_DEVICES"]):
        yticks.append(row_idx[("QC", d)])
        caps = ",".join(sorted(dev_caps[d])) if d < len(dev_caps) else "-"
        ylabels.append(f"QC-E{d+1} ({caps})")

    ax.set_yticks(yticks); ax.set_yticklabels(ylabels, color='black')
    ax.set_xlabel("Time (h)", color='black'); ax.set_title(title, color='black')
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    ax.grid(axis='x', alpha=0.3, color='black')
    for spine in ax.spines.values():
        spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()

def plot_convergence(hist, title="GA Convergence (min Cmax)"):
    if not HAS_MPL:
        print("(matplotlib not installed; skip plotting)"); return
    plt.figure(figsize=(6,4))
    plt.plot(hist, linewidth=2, color='black')
    plt.xlabel("Generation", color='black'); plt.ylabel("Best Cmax", color='black')
    plt.title(title, color='black'); plt.grid(True, alpha=0.3, color='black')
    ax=plt.gca()
    ax.tick_params(axis='x', colors='black'); ax.tick_params(axis='y', colors='black')
    for spine in ax.spines.values():
        spine.set_edgecolor('black')
    plt.tight_layout(); plt.show()


# ============================ Main ============================ #

def main():
    random.seed(CONFIG["RANDOM_SEED"]); np.random.seed(CONFIG["RANDOM_SEED"])

    # 1) instance
    bottles, dev_caps = RandomData(CONFIG).gen()

    # 2) MIX schedule
    sch = Scheduler(CONFIG, bottles, dev_caps)
    mix_end, mix_tasks = sch.schedule_mixing()

    # 3) GA over QC grouping policies
    ga = GA(CONFIG, sch, mix_end)
    best_ind, best_cmax, hist = ga.run()

    # 4) Decode best groups and schedule QC
    bottle_groups={}
    for i,b in enumerate(bottles):
        comps=[c for c,_ in b.recipe.components]
        bottle_groups[b.id] = sch.make_groups(comps, int(best_ind[i]))
    Cmax, qc_tasks = sch.schedule_qc_with_groups(mix_end, bottle_groups)

    print(f"Bottles: {CONFIG['NUM_BOTTLES']}, Mixers: {CONFIG['NUM_MIX_MACHINES']}, QC devices: {CONFIG['NUM_QC_DEVICES']}")
    print("Best grouping policy codes per bottle (0=ALL, 1=SPLIT_LONGEST, 2=EACH):")
    print(best_ind.tolist())
    print(f"Best Cmax: {Cmax:.3f} h")

    if HAS_MPL:
        plot_convergence(hist)
        plot_gantt(mix_tasks, qc_tasks, CONFIG, dev_caps)
    else:
        print("(Install matplotlib to view plots.)")

if __name__ == "__main__":
    main()
